#include "stdafx.h"
#include "DiskMgmt.h"

// Utility class to enumerate disks
class CVdsEnum
{
	IVdsService *const m_pService;
	CAtlArray<CComPtr<IVdsDisk>> &m_list;

	inline CVdsEnum( IVdsService *pService, CAtlArray<CComPtr<IVdsDisk>> &dest ) :
		m_pService( pService ),
		m_list( dest )
	{
		m_list.RemoveAll();
	}

	typedef HRESULT( CVdsEnum::*EnumProc )( IUnknown* pUnk );

	// Iterate through some IEnumVdsObject collection, call the function pointer for each item
	HRESULT vdsEnum( IEnumVdsObject *pEnum, EnumProc onEnum )
	{
		while( true )
		{
			CComPtr<IUnknown> pUnknown;
			ULONG ulFetched = 0;
			HRESULT hr = pEnum->Next( 1, &pUnknown, &ulFetched );
			CHECK( hr, "IEnumVdsObject::Next failed" );

			if( S_FALSE == hr ) // The specified number of returned objects is greater than the number of objects remaining
				return S_OK;

			if( NULL == pUnknown )
				CHECK( E_FAIL, "IEnumVdsObject::Next returned NULL" );

			hr = ( this->*onEnum )( pUnknown );
			CHECK0( hr );
		}
	}

	HRESULT onDisk( IUnknown* pUnk )
	{
		CComPtr<IVdsDisk> pDisk;
		HRESULT hr = pUnk->QueryInterface( &pDisk );
		CHECK0( hr );
		m_list.Add( pDisk );
		return S_OK;
	}

	HRESULT onPack( IUnknown* pUnk )
	{
		CComPtr<IVdsPack> pPack;
		HRESULT hr = pUnk->QueryInterface( &pPack );
		CHECK0( hr );

		/*
		// Nothing interesting here
		VDS_PACK_PROP props;
		hr = pPack->GetProperties( &props );
		CHECK0( hr );
		CString packName = props.pwszName;
		CoTaskMemFree( props.pwszName ); */

		CComPtr<IEnumVdsObject> e;
		hr = pPack->QueryDisks( &e );
		CHECK0( hr );
		hr = vdsEnum( e, &CVdsEnum::onDisk );
		CHECK( hr, "Unable to enumerate disks" );

		return S_OK;
	}

	HRESULT onSubSystem( IUnknown* pUnk )
	{
		CComPtr<IVdsSubSystem > pSubsystem;
		HRESULT hr = pUnk->QueryInterface( &pSubsystem );
		CHECK0( hr );

		return E_NOTIMPL;
	}

	HRESULT onProvider( IUnknown* pUnk )
	{
		CComPtr<IVdsHwProvider> hwProvider;
		if( SUCCEEDED( pUnk->QueryInterface( &hwProvider ) ) )
		{
			// A hardware provider (not implemented, don't have the hardware to test)
			CComPtr<IEnumVdsObject> e;
			HRESULT hr = hwProvider->QuerySubSystems( &e );
			CHECK( hr, "IVdsSwProvider::QuerySubSystems failed" );

			hr = vdsEnum( e, &CVdsEnum::onSubSystem );
			CHECK( hr, "Unable to enumerate subsystems" );

			return S_OK;
		}

		CComPtr<IVdsSwProvider > swProvider;
		if( SUCCEEDED( pUnk->QueryInterface( &swProvider ) ) )
		{
			// A software provider
			CComPtr<IEnumVdsObject> e;
			HRESULT hr = swProvider->QueryPacks( &e );
			CHECK( hr, "IVdsSwProvider::QueryPacks failed" );

			hr = vdsEnum( e, &CVdsEnum::onPack );
			CHECK( hr, "Unable to enumerate packs" );

			return S_OK;
		}
		return S_FALSE;
	}

	HRESULT list()
	{
		CComPtr<IEnumVdsObject> e;
		HRESULT hr = m_pService->QueryProviders( VDS_QUERY_HARDWARE_PROVIDERS | VDS_QUERY_SOFTWARE_PROVIDERS, &e );
		CHECK( hr, "IVdsService::QueryProviders failed" );

		hr = vdsEnum( e, &CVdsEnum::onProvider );
		CHECK( hr, "Unable to enumerate providers" );

		return S_OK;
	}

public:
	inline static HRESULT vdsList( IVdsService *pService, CAtlArray<CComPtr<IVdsDisk>> &dest )
	{
		CVdsEnum e( pService, dest );
		return e.list();
	}
};

HRESULT vdsList( IVdsService *pService, CAtlArray<CComPtr<IVdsDisk>> &dest )
{
	return CVdsEnum::vdsList( pService, dest );
}